package raft

import (
	"distributed_db/utils"
	"encoding/gob"
	"fmt"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

const Debug = 0

func DPrintf(format string, a ...interface{}) (n int, err error) {
	if Debug > 0 {
		log.Printf(format, a...)
	}
	return
}

type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
	CommitIndex int
	SerialNum   uint32
	ClientId    int64
	Operation   string
	Key         string
	Value       string
}

type OpJob struct {
	CommitIndex int
	SerialNum   uint32
	ClientId    int64
	Result      chan bool
}

func (oj OpJob) Execute(args interface{}) {
	defer close(oj.Result)
	if op, ok := args.(Op); ok {
		// raft.LogC(raft.Purple, "Triggered id: %d\n", oj.SerialNum)
		if op.ClientId == oj.ClientId && op.SerialNum == oj.SerialNum {
			oj.Result <- true
			return
		}
	}
	oj.Result <- false
}

func (oj OpJob) CondTrigger(e utils.JobEvent) bool {
	if op, ok := e.Args.(Op); ok {
		if op.CommitIndex == oj.CommitIndex {
			return true
		}
	}
	return false
}

type client struct {
	record uint32
	cond   *sync.Cond
}

type KVServer struct {
	mu      sync.Mutex
	me      int
	rf      *Raft
	applyCh chan ApplyMsg
	dead    int32 // set by Kill()

	maxraftstate int // snapshot if log grows this big

	// Your definitions here.
	kvPair       map[string]string
	clientRecord map[int64]*client
	pm           *utils.ProcessManager

	commit *sync.Cond
	Index  int // committed index

	dying *sync.Cond
	// debug
	Log []Op
}

func (kv *KVServer) get_kv_pair(key string) (string, bool) {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	v, ok := kv.kvPair[key]
	return v, ok
}

func (kv *KVServer) set_kv_pair(key string, value string) {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	kv.kvPair[key] = value
}

func (kv *KVServer) get_client_pair(key int64) (uint32, bool) {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	v, ok := kv.clientRecord[key]
	if !ok {
		return 0, ok
	}
	return v.record, ok
}

func (kv *KVServer) set_client_pair(key int64, record uint32) {
	kv.mu.Lock()
	if _, exist := kv.clientRecord[key]; !exist {
		kv.clientRecord[key] = &client{
			record: record,
			cond:   sync.NewCond(&sync.Mutex{}),
		}
	} else {
		kv.clientRecord[key].record = record
	}
	// raft.LogC(raft.Yellow, "S%d Notify %d\n", kv.me, record)
	kv.mu.Unlock()
	kv.NotifyClientValue(key)
}

func (kv *KVServer) WaitClientValue(client int64, value uint32) bool {
	kv.clientRecord[client].cond.L.Lock()
	defer kv.clientRecord[client].cond.L.Unlock()

	// 设置超时时间
	timer := time.After(time.Second)

	for value > kv.clientRecord[client].record {
		// 使用 select 实现带超时的等待
		select {
		case <-timer:
			// 超时处理逻辑
			return false
		default:
			// 等待通知
			kv.clientRecord[client].cond.Wait()
		}
	}
	return true
}

func (kv *KVServer) NotifyClientValue(client int64) {
	kv.clientRecord[client].cond.Broadcast()
}

func (kv *KVServer) get_commit_index() int {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	return kv.Index
}

func (kv *KVServer) set_commit_index(index int) {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	kv.Index = index
}

func (kv *KVServer) WaitCommitValue(index int) bool {
	kv.commit.L.Lock()
	defer kv.commit.L.Unlock()
	// 设置超时时间
	timer := time.After(time.Second)

	for index > kv.get_commit_index() {
		select {
		case <-timer:
			// 超时处理逻辑
			return false
		default:
			// 等待通知
			kv.commit.Wait()
		}
	}
	return true
}

func (kv *KVServer) NotifyCommitValue() {
	kv.commit.Broadcast()
}

func (kv *KVServer) apply_op_to_kv_pair(key string, value string, op string) (string, bool) {
	switch op {
	case "Put":
		kv.set_kv_pair(key, value)
		return "", true
	case "Append":
		if v, ok := kv.get_kv_pair(key); ok {
			kv.set_kv_pair(key, v+value)
		} else {
			kv.set_kv_pair(key, value)
		}
		return "", true
	case "Get":
		return kv.get_kv_pair(key)
	}
	return "", false
}

func (kv *KVServer) Get(args *utils.GetArgs, reply *utils.GetReply) {
	// Your code here.
	leader := kv.rf.GetLeaderId()
	// raft.LogC(raft.Green, "C%d Get S%d as entry %d\n", args.ClientId%100, kv.me, args.EntryId)
	if leader != kv.me {
		reply.Err = utils.ErrWrongLeader
		return
	}
	// check whether current entry id has been served
	v, ok := kv.get_client_pair(args.ClientId)
	if !ok {
		kv.set_client_pair(args.ClientId, args.EntryId-1)
		v = args.EntryId - 1
	}
	// if not serverd, wait for previous
	if v < args.EntryId {
		if !kv.WaitClientValue(args.ClientId, args.EntryId-1) {
			reply.Err = ""
			return
		}
		kv.set_client_pair(args.ClientId, args.EntryId)
	}
	if v, ok := kv.get_kv_pair(args.Key); ok {
		reply.Err = utils.OK
		reply.Value = v
		// raft.LogC(raft.Purple, "get at S%d, K: %v, V: %v\n", kv.me, args.Key, v)
	} else {
		reply.Err = utils.ErrNoKey
		// raft.LogC(raft.Red, "C%d Get Wrong Key as entry %d\n", args.ClientId%100, args.EntryId)
	}
}

func (kv *KVServer) PutAppend(args *utils.PutAppendArgs, reply *utils.PutAppendReply) {
	// Your code here.
	for {
		leader := kv.rf.GetLeaderId()
		if leader == kv.me {

		}
		if leader != kv.me {
			reply.Err = utils.ErrWrongLeader
			// raft.LogC(raft.Red, "C%d Get Wrong Key as entry %d\n", args.ClientId%100, args.EntryId)
			return
		}
		// Your code here.
		utils.LogC(utils.Green, "C%d Append S%d %v: %v as entry %d\n", args.ClientId%100, kv.me, args.Key, args.Value, args.EntryId)
		// check whether current entry id has been served
		v, ok := kv.get_client_pair(args.ClientId)
		if !ok {
			kv.set_client_pair(args.ClientId, args.EntryId-1)
			v = args.EntryId - 1
		}
		// if serverd, return
		if v >= args.EntryId {
			reply.Err = utils.OK
			return
		}

		// start serve current agreement

		op := Op{
			Operation: args.Op,
			Key:       args.Key,
			Value:     args.Value,
			ClientId:  args.ClientId,
			SerialNum: args.EntryId,
		}
		// start agreement
		idx, _, ok := kv.rf.Start(op)
		if !ok {
			reply.Err = utils.ErrWrongLeader
			return
		}

		judge := kv.WaitCommitValue(idx)
		if !judge {
			reply.Err = utils.ErrTimeWait
			return
		}
		v, _ = kv.get_client_pair(args.ClientId)
		if v == args.EntryId {
			// if committed as the current client and entry
			reply.Err = utils.OK
			return
		}
	}
}

// the tester calls Kill() when a KVServer instance won't
// be needed again. for your convenience, we supply
// code to set rf.dead (without needing a lock),
// and a killed() method to test rf.dead in
// long-running loops. you can also add your own
// code to Kill(). you're not required to do anything
// about this, but it may be convenient (for example)
// to suppress debug output from a Kill()ed instance.
func (kv *KVServer) Kill() {
	atomic.StoreInt32(&kv.dead, 1)
	kv.rf.Kill()
	fmt.Println("kill done")
	// Your code here, if desired.
	kv.dying.Broadcast()
}

func (kv *KVServer) killed() bool {
	z := atomic.LoadInt32(&kv.dead)
	return z == 1
}

// servers[] contains the ports of the set of
// servers that will cooperate via Raft to
// form the fault-tolerant key/value service.
// me is the index of the current server in servers[].
// the k/v server should store snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
// the k/v server should snapshot when Raft's saved state exceeds maxraftstate bytes,
// in order to allow Raft to garbage-collect its log. if maxraftstate is -1,
// you don't need to snapshot.
// StartKVServer() must return quickly, so it should start goroutines
// for any long-running work.
func StartKVServer(servers []*utils.ClientEnd, me int, persister *Persister, maxraftstate int) *KVServer {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	gob.Register(Op{})

	kv := new(KVServer)
	kv.me = me
	kv.maxraftstate = maxraftstate

	// You may need initialization code here.

	kv.applyCh = make(chan ApplyMsg)
	kv.rf = Make(servers, me, persister, kv.applyCh)
	kv.clientRecord = make(map[int64]*client)
	kv.kvPair = make(map[string]string)

	// You may need initialization code here.
	kv.pm = utils.NewProcessManager(10)
	kv.dying = sync.NewCond(&sync.Mutex{})
	kv.commit = sync.NewCond(&sync.Mutex{})
	kv.Log = make([]Op, 0)
	go func() {
		for m := range kv.applyCh {
			if !m.CommandValid {
				// ignore other types of ApplyMsg
			} else {
				// inc committed log

				v := m.Command
				op, ok := v.(Op)
				if !ok {
					utils.LogC(utils.Red, "Type assertion failed\n")
					continue
				}
				if len(kv.Log) == m.CommandIndex-1 {
					kv.Log = append(kv.Log, op)
				}
				// get client entry size
				entry, ok := kv.get_client_pair(op.ClientId)
				if !ok { // no such client
					kv.set_client_pair(op.ClientId, op.SerialNum-1)
					entry = op.SerialNum - 1
				}
				// if op num is greater then current entry size
				if entry < op.SerialNum {
					// apply op
					kv.apply_op_to_kv_pair(op.Key, op.Value, op.Operation)
					// update client serial number
					entry = op.SerialNum
					kv.set_client_pair(op.ClientId, entry)
					// raft.LogC(raft.Blue, "%s at S%d K: %v, V: %v\n", op.Operation, kv.me, op.Key, op.Value)
				}
				if m.CommandIndex > kv.get_commit_index() {
					kv.set_commit_index(m.CommandIndex)
					kv.NotifyCommitValue()
				}
			}
		}
	}()

	go func() {
		kv.dying.L.Lock()
		for !kv.killed() {
			kv.dying.Wait()
		}
		kv.dying.L.Unlock()
		close(kv.applyCh)
	}()

	return kv
}
